

from datetime import datetime
from typing import Dict, List
from pandas import DataFrame
import pandas as pd
from connector.stockConnector import stockConnector
from connector.scenarioConnector import scenarioConnector
from model.analyzer import Analyzer
from model.executor import Executor
from model.expositor import Expositor
from model.strategy import Strategy


class BacktestService:
    def __init__(self) -> None:
        self.analyzer = Analyzer()
        self.strategy_list:list[Strategy] = []
        self.date_list = []

        return

    
    def initStockBacktest(self,scenario_id,start_date,end_date,stock_id,cash,indicator_type,custom_indicator_list,strategy_description_list,item_uuid,compare_stock_id)->None:
        self.scenario_id = scenario_id
        self.executor = Executor(datetime.fromtimestamp(start_date),datetime.fromtimestamp(end_date),cash)
        self.analyzer.compare_stock['stock_code'] = stockConnector.getStockCodeById(compare_stock_id)
        self.analyzer.compare_stock['name'] = stockConnector.getStockNameByCode(self.analyzer.compare_stock['stock_code'])
        self.analyzer.start_date = self.executor.start_date
        self.analyzer.end_date = self.executor.end_date
        if indicator_type != 'custom':
            self.analyzer.indicator_type = indicator_type
        else:
            self.analyzer.indicator_type = custom_indicator_list
        self.analyzer.item_list.append(scenarioConnector.getItemByUuid(scenario_id,item_uuid))
        stock_code = stockConnector.getStockCodeById(stock_id)
        for i in strategy_description_list:
            temp_strategy_dict = Expositor().stockStrategyDscriptionParse(i['description'],stock_code)
            temp_strategy = Strategy()
            for i in temp_strategy_dict["condition_list"]:
                temp_strategy.addCondition(i['attribute_source'],i['attribute_name'],i['calculate_type'],i['variation'],i['offdays'])
            temp_strategy.setOperation(temp_strategy_dict['operation']['stock_code'],temp_strategy_dict['operation']['operation_type'],temp_strategy_dict['operation']['amount'])
            self.strategy_list.append(temp_strategy)
        return self.startBacktest()
        
        
    def initPortfolioBacktest(self,scenario_id,start_date,end_date,cash,indicator_type,custom_indicator_list,init_investment,strategy_description_list,compare_stock_id)->None:
        self.scenario_id = scenario_id
        self.executor = Executor(datetime.fromtimestamp(start_date),datetime.fromtimestamp(end_date),cash)
        self.analyzer.start_date = self.executor.start_date
        self.analyzer.end_date = self.executor.end_date
        if indicator_type != 'custom':
            self.analyzer.indicator_type = indicator_type
        else:
            self.analyzer.indicator_type = custom_indicator_list
        self.analyzer.compare_stock['stock_code'] = stockConnector.getStockCodeById(compare_stock_id)
        self.analyzer.compare_stock['name'] = stockConnector.getStockNameByCode(self.analyzer.compare_stock['stock_code'])
        # self.startBacktest(init_investment)
        for i in strategy_description_list:
            temp_strategy_dict = Expositor().portfolioStrategyDescriptionParse(i['description'])
            temp_strategy = Strategy()
            for i in temp_strategy_dict["condition_list"]:
                temp_strategy.addCondition(i['attribute_source'],i['attribute_name'],i['calculate_type'],i['variation'],i['offdays'])
            temp_strategy.setOperation(temp_strategy_dict['operation']['stock_code'],temp_strategy_dict['operation']['operation_type'],temp_strategy_dict['operation']['amount'])
            self.strategy_list.append(temp_strategy)
        
        return self.startBacktest(init_investment)
    
    

    def startBacktest(self,init_investment=[])->None:

        self.date_list = self.getDateList(init_investment)
        scenario = scenarioConnector.getScenarioById(self.scenario_id)
        for i in scenario.attribute_list:
            self.analyzer.attribute_list.append({
                "label":i['label'],
                "data_list":[]
            })
        for index,current_date in enumerate(self.date_list):
            temp_record = {
                    "date":current_date,
                    "operation_list":[],
                    "data":{}
            }
            for i in scenario.attribute_list:
                 data_list = next((item for item in self.analyzer.attribute_list if item['label'] == i['label']), None)["data_list"]
                 data_list.append(scenarioConnector.getAttributeDataOnDate(self.scenario_id,i['label'],current_date))
            if index==0:
                for i in init_investment:
                    temp_record['operation_list'].append(self.executorInitInvestmentItem(i,current_date))   
            for temp_strategy in self.strategy_list:
                if self.implStrategy(index,temp_strategy):
                   self.executor.executeOperation(temp_strategy.operation["stock_code"],temp_strategy.operation["operation_type"],self.getAttributeOfStockOnDate(temp_strategy.operation["stock_code"],current_date,"close"),temp_strategy.operation["amount"])
                   temp_record["operation_list"].append(temp_strategy.operation)
            temp_record["data"]["cash"] = self.executor.cash
            temp_record["data"]['total_assets'] = self.calTotalAssets(current_date)
            
            self.analyzer.addRecord(temp_record)
            self.analyzer.compare_stock['data'].append(self.getAttributeOfStockOnDate(self.analyzer.compare_stock['stock_code'],current_date,'close'))
        report = self.analyzer.getBacktestReport()
        for i in report['operation_list']:
            for j in i:
                j['name'] = stockConnector.getStockNameByCode(j['stock_code'])
        return report


    def calTotalAssets(self,date)->float:
        total_assets = 0
        for i in self.executor.position:
            stock_value = self.getAttributeOfStockOnDate(i['stock_code'],date,"close")*i['amount']
            total_assets += stock_value
        total_assets += self.executor.cash
        return total_assets
    
    def getDateList(self,init_investment=[])->List[datetime]:
        dfs = []
        for i in self.strategy_list:
            for j in i.condition_list:
                dfs.append(self.getAttributeData(j["attribute_source"],j["attribute_name"],self.executor.start_date,self.executor.end_date))
        for i in init_investment:
            dfs.append(self.getAttributeData(stockConnector.getStockCodeById(i["stock_id"]),"close",self.executor.start_date,self.executor.end_date))
        # 初始化combined_df为第一个DataFrame
        combined_df = dfs[0]
        # 使用for循环，基于 'key' 列合并剩余的DataFrame
        for df in dfs[1:]:
            combined_df = pd.merge(combined_df, df, on='date', how='inner')
        return  combined_df["date"].tolist()
    
    def getAttributeData(self,attribute_source,attribute_name,start_date,end_date)->DataFrame:
            if attribute_source == "scenario":
                return scenarioConnector.getAttributeDataInRange(self.scenario_id,attribute_name,start_date,end_date)
            else:
                return stockConnector.getAttributeDataInRange(attribute_source,attribute_name,start_date,end_date)
            
    def implStrategy(self,curr_index:int,strategy:Strategy)->None:
        for temp_condition in strategy.condition_list:
            if temp_condition["offdays"] > curr_index:
                return False
            data = self.getAttributeData(temp_condition["attribute_source"],temp_condition["attribute_name"],self.date_list[curr_index-temp_condition["offdays"]],self.date_list[curr_index])
            if self.calculateCondition(data,temp_condition["calculate_type"],temp_condition["variation"]) == False:
                return False
        return True
    
    def getAttributeOfStockOnDate(self,stock_code:str,date:datetime,attribute_name:str)->float:
        return stockConnector.getAttributeOfStockOnDate(stock_code,date,attribute_name)
    
    def getItemListByDescription(self,scenario_id,description)->list:
        scenario = scenarioConnector.getScenarioById(scenario_id)
        description_list = description.split('/')
        condition_list = []
        for i in description_list:
            condition_list.append(Expositor().itemDescriptionParse(i))
        temp_item_list = scenario.item_list
        for i in temp_item_list:
            score = 0
            for j in i['attribute_list']:
                for k in condition_list:
                    if j['label'] == k['label']:
                        if j['value'] == k['value']:
                            score += 1
            i['score']=score
        sorted_item_list = sorted(temp_item_list,key=lambda x:x['score'],reverse=True)[:5]
        return sorted_item_list
    
    def executorInitInvestmentItem(self,init_investment_item,date)->None:
        temp_cash = float(self.executor.cash)*int(init_investment_item["value"])/100
        stock_code = stockConnector.getStockCodeById(init_investment_item["stock_id"])
        price = self.getAttributeOfStockOnDate(stock_code,date,"close")
        amount = (int((temp_cash/price)/100))*100
        self.executor.executeOperation(stock_code,"BUY",price,amount)
        return {
            "stock_code":stock_code,
            "operation_type":'BUY',
            "amount":amount
        }

    def calculateCondition(self,data:DataFrame,calculate_type:str,variation:float)->bool:
        if calculate_type == 'avarage':
            sum = data.iloc[:,1].sum()
            avarage_chg = sum/(data.shape[0]-1)
            return avarage_chg>=variation
        
        if calculate_type == 'rise':
            chg = (data.iloc[-1,1] - data.iloc[0,1])/data.iloc[0,1]
            return (chg*100)>=variation
        if calculate_type == 'fall':
            chg = (data.iloc[-1,1] - data.iloc[0,1])/data.iloc[0,1]
            return (chg*100)<=variation